#include <node.h>
#include <node_buffer.h>  // 添加 Buffer 相关头文件
#include <windows.h>
#include <gdiplus.h>
#include <vector>

using namespace v8;
using namespace Gdiplus;

#pragma comment(lib, "gdiplus.lib")

// 获取 JPEG 编码器的 CLSID
int GetEncoderClsid(const WCHAR* format, CLSID* pClsid) {
  UINT num = 0;
  UINT size = 0;
  GetImageEncodersSize(&num, &size);
  if (size == 0) return -1;
  
  ImageCodecInfo* pImageCodecInfo = (ImageCodecInfo*)malloc(size);
  GetImageEncoders(num, size, pImageCodecInfo);
  
  for (UINT j = 0; j < num; ++j) {
    if (wcscmp(pImageCodecInfo[j].MimeType, format) == 0) {
      *pClsid = pImageCodecInfo[j].Clsid;
      free(pImageCodecInfo);
      return j;
    }
  }
  
  free(pImageCodecInfo);
  return -1;
}

void CaptureScreen(const FunctionCallbackInfo<Value>& args) {
  Isolate* isolate = args.GetIsolate();

  // 1. GDI+ 初始化
  GdiplusStartupInput gdiplusStartupInput;
  ULONG_PTR gdiplusToken;
  if (GdiplusStartup(&gdiplusToken, &gdiplusStartupInput, NULL) != Ok) {
    return;
  }

  // 2. 获取屏幕 DC
  HDC hdcScreen = GetDC(NULL);
  if (!hdcScreen) {
    GdiplusShutdown(gdiplusToken);
    return;
  }

  // 3. 创建内存 DC
  HDC hdcMem = CreateCompatibleDC(hdcScreen);
  if (!hdcMem) {
    ReleaseDC(NULL, hdcScreen);
    GdiplusShutdown(gdiplusToken);
    return;
  }

  // 4. 获取屏幕尺寸
  int width = GetSystemMetrics(SM_CXSCREEN);
  int height = GetSystemMetrics(SM_CYSCREEN);

  // 5. 创建位图
  HBITMAP hBitmap = CreateCompatibleBitmap(hdcScreen, width, height);
  if (!hBitmap) {
    DeleteDC(hdcMem);
    ReleaseDC(NULL, hdcScreen);
    GdiplusShutdown(gdiplusToken);
    return;
  }

  // 6. 选入位图
  SelectObject(hdcMem, hBitmap);

  // 7. 屏幕拷贝
  if (!BitBlt(hdcMem, 0, 0, width, height, hdcScreen, 0, 0, SRCCOPY)) {
    DeleteObject(hBitmap);
    DeleteDC(hdcMem);
    ReleaseDC(NULL, hdcScreen);
    GdiplusShutdown(gdiplusToken);
    return;
  }

  IStream* stream = NULL;
  if (CreateStreamOnHGlobal(NULL, TRUE, &stream) != S_OK) {
    DeleteObject(hBitmap);
    DeleteDC(hdcMem);
    ReleaseDC(NULL, hdcScreen);
    GdiplusShutdown(gdiplusToken);
    return;
  }

  CLSID clsid;
  if (GetEncoderClsid(L"image/png", &clsid) == -1) {
    stream->Release();
    DeleteObject(hBitmap);
    DeleteDC(hdcMem);
    ReleaseDC(NULL, hdcScreen);
    GdiplusShutdown(gdiplusToken);
    return;
  }

  Bitmap* pBitmap = Bitmap::FromHBITMAP(hBitmap, NULL);
  if (!pBitmap || pBitmap->GetLastStatus() != Ok) {
    stream->Release();
    DeleteObject(hBitmap);
    DeleteDC(hdcMem);
    ReleaseDC(NULL, hdcScreen);
    GdiplusShutdown(gdiplusToken);
    return;
  }
  Status saveStatus = pBitmap->Save(stream, &clsid);
  delete pBitmap;
  if (saveStatus != Ok) {
    stream->Release();
    DeleteObject(hBitmap);
    DeleteDC(hdcMem);
    ReleaseDC(NULL, hdcScreen);
    GdiplusShutdown(gdiplusToken);
    return;
  }

  LARGE_INTEGER liZero = {};
  stream->Seek(liZero, STREAM_SEEK_SET, NULL);

  STATSTG stat;
  if (stream->Stat(&stat, STATFLAG_NONAME) != S_OK) {
    stream->Release();
    DeleteObject(hBitmap);
    DeleteDC(hdcMem);
    ReleaseDC(NULL, hdcScreen);
    GdiplusShutdown(gdiplusToken);
    return;
  }

  HGLOBAL hMem = NULL;
  if (GetHGlobalFromStream(stream, &hMem) != S_OK) {
    stream->Release();
    DeleteObject(hBitmap);
    DeleteDC(hdcMem);
    ReleaseDC(NULL, hdcScreen);
    GdiplusShutdown(gdiplusToken);
    return;
  }
  LPVOID pvData = GlobalLock(hMem);
  if (!pvData) {
    stream->Release();
    DeleteObject(hBitmap);
    DeleteDC(hdcMem);
    ReleaseDC(NULL, hdcScreen);
    GdiplusShutdown(gdiplusToken);
    return;
  }

  size_t dataLen = stat.cbSize.LowPart;
  std::vector<char> data((char*)pvData, (char*)pvData + dataLen);
  MaybeLocal<Object> buffer = node::Buffer::Copy(isolate, data.data(), data.size());
  if (!buffer.IsEmpty()) {
    args.GetReturnValue().Set(buffer.ToLocalChecked());
  }

  GlobalUnlock(hMem);
  stream->Release();
  DeleteObject(hBitmap);
  DeleteDC(hdcMem);
  ReleaseDC(NULL, hdcScreen);
  GdiplusShutdown(gdiplusToken);
}

void Initialize(Local<Object> exports) {
  NODE_SET_METHOD(exports, "capture", CaptureScreen);
}

NODE_MODULE(NODE_GYP_MODULE_NAME, Initialize)